import pathlib

import scipy
import torch
import torch.utils.data


class KdVDataset(torch.utils.data.Dataset):
    def __init__(self,
                 path: pathlib.Path = pathlib.Path("./data/KdV.mat")):
        data = scipy.io.loadmat(str(path))

        t = torch.tensor(data['tt'], dtype=torch.float)
        x = torch.tensor(data['x'], dtype=torch.float)
        u = torch.tensor(data['uu'], dtype=torch.float)

        t = torch.flatten(t)
        x = torch.flatten(x)

        t, x = torch.meshgrid(t, x, indexing="xy")

        t = torch.flatten(t)
        x = torch.flatten(x)
        u = torch.flatten(u)

        self._data = t, x, u

    def __getitem__(self, item):
        t, x, u = self._data
        return t[item], x[item], u[item]

    def __len__(self):
        return self._data[0].shape[0]
